import os
import shutil
import random
 
def make_yolo_dataset(images_folder, labels_folder, output_folder, train_ratio=0.8):
    # 创建目标文件夹
    images_train_folder = os.path.join(output_folder, '../data/train/images')
    images_val_folder = os.path.join(output_folder, '../data/valid/images')
    labels_train_folder = os.path.join(output_folder, '../data/train/labels')
    labels_val_folder = os.path.join(output_folder, '../data/valid/labels')
 
    os.makedirs(images_train_folder, exist_ok=True)
    os.makedirs(images_val_folder, exist_ok=True)
    os.makedirs(labels_train_folder, exist_ok=True)
    os.makedirs(labels_val_folder, exist_ok=True)
 
    # 获取图片和标签的文件名（不包含扩展名）
    image_files = [f for f in os.listdir(images_folder) if f.endswith('.jpg')]
    label_files = [f for f in os.listdir(labels_folder) if f.endswith('.txt')]
    image_base_names = set(os.path.splitext(f)[0] for f in image_files)
    label_base_names = set(os.path.splitext(f)[0] for f in label_files)
 
    # 找出图片和标签都存在的文件名
    matched_files = list(image_base_names & label_base_names)
 
    # 打乱顺序并划分为训练集和验证集
    random.shuffle(matched_files)
    split_idx = int(len(matched_files) * train_ratio)
    train_files = matched_files[:split_idx]
    val_files = matched_files[split_idx:]
 
    # 移动文件到对应文件夹
    for base_name in train_files:
        img_src = os.path.join(images_folder, f"{base_name}.jpg")
        lbl_src = os.path.join(labels_folder, f"{base_name}.txt")
 
        img_dst = os.path.join(images_train_folder, f"{base_name}.jpg")
        lbl_dst = os.path.join(labels_train_folder, f"{base_name}.txt")
 
        shutil.copyfile(img_src, img_dst)
        shutil.copyfile(lbl_src, lbl_dst)
 
    for base_name in val_files:
        img_src = os.path.join(images_folder, f"{base_name}.jpg")
        lbl_src = os.path.join(labels_folder, f"{base_name}.txt")
 
        img_dst = os.path.join(images_val_folder, f"{base_name}.jpg")
        lbl_dst = os.path.join(labels_val_folder, f"{base_name}.txt")
 
        shutil.copyfile(img_src, img_dst)
        shutil.copyfile(lbl_src, lbl_dst)
 
    print("数据集划分完成！")
 

if __name__ == '__main__' :
    images_folder = '../images'  # 原始图片文件夹路径
    labels_folder = '../labels'  # 原始标签文件夹路径
    output_folder = '/'  # 存放结果数据集的文件夹路径
    make_yolo_dataset(images_folder, labels_folder, output_folder)
 